"""
Routing visualization using Mermaid diagrams
"""
from typing import List, Dict, Set
from .base import MessageRouter
from .conditions import ConditionalRouter
from .broadcast import BroadcastRouter

class RoutingVisualizer:
    """路由可视化器"""
    
    def __init__(self, message_router: MessageRouter,
                 conditional_router: ConditionalRouter,
                 broadcast_router: BroadcastRouter):
        """
        初始化可视化器
        Args:
            message_router: 消息路由器实例
            conditional_router: 条件路由器实例
            broadcast_router: 广播路由器实例
        """
        self.message_router = message_router
        self.conditional_router = conditional_router
        self.broadcast_router = broadcast_router
        
    def generate_mermaid(self, include_conditions: bool = True,
                        include_broadcasts: bool = True) -> str:
        """
        生成完整的Mermaid图表
        Args:
            include_conditions: 是否包含条件路由
            include_broadcasts: 是否包含广播规则
        Returns:
            Mermaid图表字符串
        """
        nodes = set()  # 所有节点集合
        edges = []     # 所有边的列表
        subgraphs = [] # 子图列表
        
        # 生成基本路由
        base_edges = self._generate_base_routes()
        edges.extend(base_edges)
        
        # 生成条件路由
        if include_conditions:
            condition_subgraph, condition_edges = self._generate_conditional_routes()
            if condition_subgraph:
                subgraphs.append(condition_subgraph)
            edges.extend(condition_edges)
            
        # 生成广播路由
        if include_broadcasts:
            broadcast_subgraph, broadcast_edges = self._generate_broadcast_routes()
            if broadcast_subgraph:
                subgraphs.append(broadcast_subgraph)
            edges.extend(broadcast_edges)
            
        # 收集所有节点
        for edge in edges:
            nodes.add(edge["from"])
            nodes.add(edge["to"])
            
        # 生成Mermaid图表
        return self._build_mermaid_diagram(nodes, edges, subgraphs)
        
    def _generate_base_routes(self) -> List[Dict]:
        """生成基本路由的边"""
        edges = []
        
        # 处理点对点路由
        for source, target in self.message_router.base_rule.routes.items():
            edges.append({
                "from": source,
                "to": target,
                "style": "normal",
                "label": ""
            })
            
        # 处理默认顺序
        order = self.message_router.base_rule.default_order
        for i in range(len(order) - 1):
            if order[i] not in self.message_router.base_rule.routes:
                edges.append({
                    "from": order[i],
                    "to": order[i + 1],
                    "style": "default_order",
                    "label": "默认"
                })
                
        return edges
        
    def _generate_conditional_routes(self) -> tuple[str, List[Dict]]:
        """生成条件路由的子图和边"""
        edges = []
        conditions_count = 0
        
        for source, conditions in self.conditional_router.conditions.items():
            for condition in conditions:
                conditions_count += 1
                edges.append({
                    "from": source,
                    "to": condition.target,
                    "style": "conditional",
                    "label": condition.description
                })
                
        if conditions_count > 0:
            subgraph = """
    subgraph 条件路由
    style 条件路由 fill:#f9f,stroke:#333,stroke-width:2px
    end
            """.strip()
        else:
            subgraph = ""
            
        return subgraph, edges
        
    def _generate_broadcast_routes(self) -> tuple[str, List[Dict]]:
        """生成广播路由的子图和边"""
        edges = []
        broadcasts_count = 0
        
        for source, patterns in self.broadcast_router.patterns.items():
            for pattern in patterns:
                broadcasts_count += 1
                for target in pattern.get_targets():
                    edges.append({
                        "from": source,
                        "to": target,
                        "style": "broadcast",
                        "label": pattern.description
                    })
                    
        if broadcasts_count > 0:
            subgraph = """
    subgraph 广播规则
    style 广播规则 fill:#ff9,stroke:#333,stroke-width:2px
    end
            """.strip()
        else:
            subgraph = ""
            
        return subgraph, edges
        
    def _build_mermaid_diagram(self, nodes: Set[str],
                              edges: List[Dict],
                              subgraphs: List[str]) -> str:
        """构建完整的Mermaid图表"""
        # 图表头部
        lines = ["graph TD", "    %% 节点样式"]
        
        # 添加节点样式
        for node in nodes:
            lines.append(f"    style {node} fill:#fff,stroke:#333,stroke-width:2px")
            
        # 添加连接线样式定义
        lines.extend([
            "    %% 连接线样式",
            "    classDef default_order stroke:#666,stroke-width:1px",
            "    classDef conditional stroke:#f66,stroke-width:2px",
            "    classDef broadcast stroke:#66f,stroke-dasharray: 5 5"
        ])
        
        # 添加边
        lines.append("    %% 连接关系")
        for edge in edges:
            style = {
                "normal": "-->",
                "default_order": "--->",
                "conditional": "==>",
                "broadcast": "-.->",
            }[edge["style"]]
            
            label = f"|{edge['label']}|" if edge["label"] else ""
            lines.append(f"    {edge['from']} {style}{label} {edge['to']}")
            
        # 添加子图
        if subgraphs:
            lines.append("    %% 子图")
            lines.extend(subgraphs)
            
        # 添加图例
        lines.extend([
            "    %% 图例",
            "    subgraph 图例",
            "    A1[正常路由] --> B1",
            "    A2[默认顺序] ---> B2",
            "    A3[条件路由] ==> B3",
            "    A4[广播规则] -.-> B4",
            "    end"
        ])
        
        return "\n".join(lines)
